# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Use grids to calculate neighbour list
"""

from typing import Tuple
import itertools
import numpy as np
import scipy.stats
import torch
from torch import Tensor
from torch.nn import Module, Parameter

from ..function.functions import get_integer, get_tensor, coordinate_in_pbc, reduce_all


class GridNeighbours(Module):
    r"""Neighbour list calculated by grids

    Args:

        cutoff (float):         Cutoff distance.

        coordinate (Tensor):    Tensor of shape `(B, A, D)`. Data type is float32.
                                position coordinates of atoms in the simulation system.

        pbc_box (Tensor):       Tensor of shape `(B, A, D)`. Data type is float32.
                                Box size of periodic boundary condition. Default: None

        atom_mask (Tensor):     Tensor of shape `(B, A)`. Data type is bool_.
                                Mask of atoms in the system.
                                Default: None

        exclude_index (Tensor): Tensor of shape `(B, A, Ex)`. Data type is int32.
                                Index of neighbour atoms which could be excluded from the neighbour list.
                                Default: None

        num_neighbours (int):   Number of neighbours. If `None` is given, this value will be calculated by
                                the ratio of the number of neighbouring grids to the total number of grids.
                                Default: None

        cell_capacity (int):    Capacity number of atoms in grid cell. If `None` is given, this value will be multiplied
                                by a factor of the maximum number of atoms in the grid cell at the initial coordinate.
                                Default: None

        num_cell_cut (int):     Number of subdivision of grid cells according to the cutoff. Default: 1

        cutoff_scale (float):   Factor to scale the cutoff distance. Default: 1.2

        cell_cap_scale (float): Factor to scale `cell_capacity`. Default: 1.25

        grid_num_scale (float): Scale factor to calculate `num_neighbours` by the ratio of grids.
                                If `num_neighbours` is not None, it will not be used. Default: 1.5

    Supported Platforms:

        ``Ascend`` ``GPU``

    Symbols:

        B:  Number of simulation walker.

        A:  Number of atoms in system.

        N:  Number of neighbour atoms.

        D:  Dimension of position coordinates.

        Ex: Maximum number of excluded neighbour atoms.

    """

    def __init__(self,
                 cutoff: float,
                 coordinate: Tensor,
                 pbc_box: Tensor = None,
                 atom_mask: Tensor = None,
                 exclude_index: Tensor = None,
                 num_neighbours: int = None,
                 cell_capacity: int = None,
                 num_cell_cut: int = 1,
                 cutoff_scale: float = 1.2,
                 cell_cap_scale: float = 1.25,
                 grid_num_scale: float = 1.5,
                 ):

        super().__init__()

        self.num_atoms = coordinate.shape[-2]
        self.dim = coordinate.shape[-1]

        self.cutoff = get_tensor(cutoff, dtype=torch.float32)

        self.cutoff_scale = get_tensor(cutoff_scale, dtype=torch.float32)
        self.cell_cap_scale = get_tensor(cell_cap_scale, dtype=torch.float32)
        self.grid_num_scale = get_tensor(grid_num_scale, dtype=torch.float32)

        # N_c
        num_cell_cut = get_integer(num_cell_cut)

        self.grid_cutoff = self.cutoff / num_cell_cut
        self.scaled_cutoff = self.cutoff * self.cutoff_scale
        self.scaled_grid_cutoff = self.grid_cutoff * self.cutoff_scale

        if pbc_box is None:
            self.use_pbc = False
            # (B, 1, D) <- (B, A, D)
            rmax = torch.max(coordinate, dim=-2, keepdim=True)
            rmin = torch.min(coordinate, dim=-2, keepdim=True)
            center = torch.mean(coordinate, dim=-2, keepdim=True)
            # (B, 2, D)
            rhalf = torch.cat((rmax-center, center-rmin), dim=-2)
            # (B, D)
            rhalf = torch.max(rhalf, dim=-2)
            # (D)
            rhalf = torch.max(rhalf, dim=0)
            box = rhalf * 2
            self.origin_grid_dims = torch.ceil(box/self.scaled_grid_cutoff).to(torch.int32)
            self.grid_dims = self.origin_grid_dims + 2
            box = self.origin_grid_dims * self.scaled_grid_cutoff
            self.half_box = box / 2
        else:
            self.use_pbc = True
            center = None
            # (B, D)
            box = get_tensor(pbc_box, dtype=torch.float32)
            if box.ndim == 1:
                box = box.unsqueeze(0)
            self.half_box = box / 2
            if (self.cutoff > self.half_box).any():
                raise ValueError(f'The cutoff ({self.cutoff}) cannot be greater than '
                                 f'the half of the length of the shortest side of '
                                 f'the PBC pbc_box {self.half_box}!')
            # (B, D)
            self.origin_grid_dims = torch.floor(box/self.scaled_grid_cutoff)
            # (D)
            self.origin_grid_dims = get_tensor(np.min(self.origin_grid_dims.cpu().numpy(), axis=0).astype(np.int32))
            self.grid_dims = self.origin_grid_dims

        # (D)
        grid_mask = self.grid_dims > (2 * num_cell_cut + 1)
        self.grid_dims = torch.where(grid_mask, self.grid_dims, torch.ones_like(self.grid_dims))
        self.max_grid_index = self.origin_grid_dims - 1

        # G
        self.num_grids = int(np.prod(self.grid_dims.detach().cpu().numpy()))

        # (D)
        self.grid_factor = torch.cumprod(self.grid_dims.flip(-1), dim=-1)
        self.grid_factor = torch.cat((self.grid_factor[:-1].flip(-1), get_tensor([1], dtype=torch.int32)), dim=-1)

        # (G, D)
        grids = [np.arange(dim).tolist() for dim in self.grid_dims.detach().cpu().numpy()]
        grids = get_tensor(tuple(itertools.product(*grids)), dtype=torch.int32)

        # (B, 1, D)
        box = box.unsqueeze(-2)
        if self.use_pbc:
            # (B, 1, D) = (B, D) / (D)
            self.cell = box / self.grid_dims
            if (self.cell < self.grid_cutoff).any():
                raise ValueError('The cell length of cannot be smaller than cutoff!')
            # (B, A, D) = ((B, A, D) - (D)) / (B, 1, D)
            atom_grid_idx = torch.floor(coordinate_in_pbc(coordinate, pbc_box)/self.cell).to(torch.int32)
        else:
            self.cell = torch.broadcast_to(self.scaled_grid_cutoff, (self.dim,))
            # (B, A, D) = (B, A, D) - (B, 1, D) + (D)
            scaled_coord = (coordinate - center + self.half_box) / self.scaled_grid_cutoff
            scaled_coord = torch.where(scaled_coord < 0, torch.zeros_like(scaled_coord), scaled_coord)
            atom_grid_idx = torch.floor(scaled_coord).to(torch.int32)
            atom_grid_idx = torch.where(atom_grid_idx < self.origin_grid_dims,
                                       atom_grid_idx, self.max_grid_index)
            atom_grid_idx += 1

        # (B, A) <- (B, A, D) * (D)
        atom_grid_idx = torch.sum(atom_grid_idx * self.grid_factor, dim=-1)

        # (B, G)
        grid_atom_num = torch.zeros((coordinate.shape[0], self.num_grids), dtype=torch.int32)
        # (B, G, C)
        if cell_capacity is None:
            # (B, G)
            grid_atom_num = torch.zeros((coordinate.shape[0], self.num_grids), dtype=torch.int32)
            for b in range(coordinate.shape[0]):
                grid_atom_num[b] = torch.bincount(atom_grid_idx[b], minlength=self.num_grids)
            cell_capacity = int(torch.max(grid_atom_num).item() * self.cell_cap_scale)

        self.cell_capacity = cell_capacity
        self.grid_atoms = torch.zeros((coordinate.shape[0], self.num_grids, self.cell_capacity), dtype=torch.int32)
        self.grid_atom_num = torch.zeros((coordinate.shape[0], self.num_grids), dtype=torch.int32)

        # (B, G, C)
        for b in range(coordinate.shape[0]):
            for i in range(self.num_atoms):
                grid_idx = atom_grid_idx[b, i]
                atom_num = self.grid_atom_num[b, grid_idx]
                if atom_num < self.cell_capacity:
                    self.grid_atoms[b, grid_idx, atom_num] = i
                    self.grid_atom_num[b, grid_idx] += 1

        # (G, D)
        self.grid_neighbours = []
        for grid in grids:
            # (D)
            grid_neighbours = []
            for d in range(self.dim):
                # (D)
                grid_d = grid.clone()
                if self.use_pbc:
                    grid_d[d] = (grid_d[d] + 1) % self.grid_dims[d]
                else:
                    grid_d[d] = grid_d[d] + 1
                    if grid_d[d] >= self.grid_dims[d]:
                        grid_d[d] = grid_d[d] - 2
                grid_neighbours.append(grid_d)
            # (D, D)
            grid_neighbours = torch.stack(grid_neighbours)
            # (D) <- (D, D) * (D)
            grid_neighbours = torch.sum(grid_neighbours * self.grid_factor, dim=-1)
            self.grid_neighbours.append(grid_neighbours)
        # (G, D)
        self.grid_neighbours = torch.stack(self.grid_neighbours)

        # (G)
        grid_idx = torch.sum(grids * self.grid_factor, dim=-1)
        # (G, D)
        self.grid_neighbours = torch.gather(self.grid_neighbours, 0, grid_idx.unsqueeze(-1).expand(-1, self.dim))

        if num_neighbours is None:
            num_neighbours = int(self.num_atoms * self.grid_num_scale)
        self.num_neighbours = num_neighbours

        self.atom_mask = atom_mask
        if atom_mask is not None:
            self.atom_mask = torch.tensor(atom_mask, dtype=torch.bool)
            if self.atom_mask.ndim == 1:
                self.atom_mask = self.atom_mask.unsqueeze(0)

        self.exclude_index = exclude_index
        if exclude_index is not None:
            self.exclude_index = torch.tensor(exclude_index, dtype=torch.int32)
            if self.exclude_index.ndim == 2:
                self.exclude_index = self.exclude_index.unsqueeze(0)

    def set_exclude_index(self, exclude_index: Tensor) -> Tensor:
        """set exclude index"""
        if exclude_index is None:
            return self
        self.exclude_index = exclude_index
        if self.exclude_index.ndim == 2:
            self.exclude_index = self.exclude_index.unsqueeze(0)
        return self.exclude_index

    def check_neighbour_list(self, raise_error: bool = True) -> bool:
        # """check the number of neighbouring atoms in neighbour list"""
        # if self.num_neighbours is not None and self.num_neighbours > self.num_atoms:
        #     if raise_error:
        #         raise RuntimeError(f'The max number of neighbour atoms ({self.num_neighbours}) '
        #                            f'is larger than the number of atoms ({self.num_atoms}!')
        #     print('[WARNING] The max number of neighbour atoms exceeds '
        #           'the initial neighbouring number of neighbour list!')
        #     return False
        return True

    def print_info(self):
        """print information of neighbour list"""
        print('[MindSPONGE] Neighbour list: GridNeighbours')
        print(f'[MindSPONGE]     Cut-off distance: {self.cutoff}')
        print(f'[MindSPONGE]     Scaled cut-off: {self.scaled_cutoff}')
        print(f'[MindSPONGE]     Grid cut-off: {self.grid_cutoff}')
        print(f'[MindSPONGE]     Scaled grid cut-off: {self.scaled_grid_cutoff}')
        print(f'[MindSPONGE]     Grid dimensions: {self.grid_dims}')
        print(f'[MindSPONGE]     Number of grids: {self.num_grids}')
        print(f'[MindSPONGE]     Cell capacity: {self.cell_capacity}')
        print(f'[MindSPONGE]     Max number of neighbour atoms: {self.num_neighbours}')
        return self

    def get_neighbours_from_grids(self, atom_grid_idx: Tensor, num_neighbours: int):
        """get neighbours from grids"""
        # (B, A, G)
        grid_idx = atom_grid_idx.unsqueeze(-1)
        grid_idx = torch.cat((grid_idx, torch.gather(self.grid_neighbours, 0,
                                                    grid_idx.expand(-1, -1, self.dim))), dim=-1)

        # (B, A, G, C)
        grid_atoms = torch.gather(self.grid_atoms, 1,
                                grid_idx.unsqueeze(-1).expand(-1, -1, -1, self.cell_capacity))
        # (B, A, G*C)
        grid_atoms = grid_atoms.reshape(grid_atoms.shape[:-2] + (-1,))

        # (B, A, G)
        grid_atom_num = torch.gather(self.grid_atom_num, 1, grid_idx)
        # (B, A)
        total_atom_num = torch.sum(grid_atom_num, dim=-1)

        # (B, A, N)
        neighbours = torch.zeros(grid_atoms.shape[:-1] + (num_neighbours,), dtype=torch.int32)
        # (B, A, N)
        neighbour_mask = torch.zeros_like(neighbours, dtype=torch.bool)

        # (B, A)
        for b in range(grid_atoms.shape[0]):
            for i in range(grid_atoms.shape[1]):
                atom_num = total_atom_num[b, i]
                if atom_num > 0:
                    atoms = grid_atoms[b, i, :atom_num]
                    atoms = atoms[atoms != i]
                    atom_num = len(atoms)
                    if atom_num > 0:
                        if atom_num > num_neighbours:
                            atoms = atoms[:num_neighbours]
                            atom_num = num_neighbours
                        neighbours[b, i, :atom_num] = atoms
                        neighbour_mask[b, i, :atom_num] = True

        return neighbours, neighbour_mask

    def forward(self,
                  coordinate: Tensor,
                  pbc_box: Tensor = None,
                  atom_mask: Tensor = None,
                  exclude_index: Tensor = None,
                  ) -> Tuple[Tensor, Tensor]:
        r"""Calculate neighbour list by grids.

        Args:
            coordinate (Tensor):    Tensor of shape (B, A, D). Data type is float.
                                    Position coordinate.
            pbc_box (Tensor):       Tensor of shape (B, D). Data type is float.
                                    Size of PBC box.
            atom_mask (Tensor):     Tensor of shape (B, A). Data type is bool.
                                    Mask of atoms.
            exclude_index (Tensor): Tensor of shape (B, A, Ex). Data type is int.
                                    Index of excluded atoms.

        Returns:
            neighbours (Tensor):     Tensor of shape (B, A, N). Data type is int.
                                    Index of neighbouring atoms.
            neighbour_mask (Tensor): Tensor of shape (B, A, N). Data type is bool.
                                    Mask of neighbour list.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
            A:  Number of atoms.
            N:  Number of neighbour atoms.
            D:  Dimension of position coordinates.
            Ex: Maximum number of excluded neighbour atoms.

        """

        if self.use_pbc:
            # (B, A, D) = ((B, A, D) - (D)) / (B, 1, D)
            atom_grid_idx = torch.floor(coordinate_in_pbc(coordinate, pbc_box)/self.cell).to(torch.int32)
        else:
            # (B, A, D) = (B, A, D) - (B, 1, D) + (D)
            center = torch.mean(coordinate, dim=-2, keepdim=True)
            scaled_coord = (coordinate - center + self.half_box) / self.scaled_grid_cutoff
            scaled_coord = torch.where(scaled_coord < 0, torch.zeros_like(scaled_coord), scaled_coord)
            atom_grid_idx = torch.floor(scaled_coord).to(torch.int32)
            atom_grid_idx = torch.where(atom_grid_idx < self.origin_grid_dims,
                                       atom_grid_idx, self.max_grid_index)
            atom_grid_idx += 1

        # (B, A) <- (B, A, D) * (D)
        atom_grid_idx = torch.sum(atom_grid_idx * self.grid_factor, dim=-1)

        # (B, A, N), (B, A, N)
        neighbours, neighbour_mask = self.get_neighbours_from_grids(atom_grid_idx, self.num_neighbours)

        if atom_mask is not None:
            if atom_mask.ndim == 1:
                atom_mask = atom_mask.unsqueeze(0)
            # (B, A, N) = (B, A, 1) & (B, A, N)
            neighbour_mask = atom_mask.unsqueeze(-1) & neighbour_mask
            # (B, A, N) = (B, 1, A) & (B, A, N)
            neighbour_mask = atom_mask.unsqueeze(1) & neighbour_mask

        if exclude_index is not None:
            if exclude_index.ndim == 2:
                exclude_index = exclude_index.unsqueeze(0)
            # (B, A, N, Ex) = (B, A, N, 1) == (B, A, 1, Ex)
            exclude_mask = neighbours.unsqueeze(-1) == exclude_index.unsqueeze(-2)
            # (B, A, N)
            exclude_mask = torch.any(exclude_mask, dim=-1)
            neighbour_mask = neighbour_mask & ~exclude_mask

        return neighbours, neighbour_mask
